import pyaudio
import numpy as np
import threading
import time
from collections import deque
import wave

class ApplicationSoundCapture:
    def __init__(self):
        self.audio = pyaudio.PyAudio()
        self.stream = None
        self.is_recording = False
        self.audio_buffer = deque(maxlen=10000)
        self.volume_threshold = 0.01
        self.callback_func = None
        self.volume=0
        
        # 录音设置
        self.sample_rate = 44100
        self.chunk_size = 1024
        self.channels = 2
        
    def list_loopback_devices(self):
        """列出环回设备（用于捕获系统播放的声音）"""
        print("=== 系统音频设备（包含环回设备）===")
        loopback_devices = []
        
        device_count = self.audio.get_device_count()
        for i in range(device_count):
            try:
                info = self.audio.get_device_info_by_index(i)
                
                # 检查是否为输出设备或包含"立体声混音"等关键词
                device_name = info['name'].lower()
                is_loopback = (
                    '立体声混音' in device_name or
                    'stereo mix' in device_name or
                    'what u hear' in device_name or
                    'wave out mix' in device_name or
                    'speakers' in device_name or
                    'headphones' in device_name
                )
                
                if info['maxInputChannels'] > 0 and is_loopback:
                    loopback_devices.append((i, info))
                    print(f"🔊 设备 {i}: {info['name']}")
                    print(f"   输入通道: {info['maxInputChannels']}")
                    print(f"   采样率: {info['defaultSampleRate']}")
                    print("-" * 40)
                
            except Exception as e:
                continue
        
        if not loopback_devices:
            print("❌ 没有找到环回设备")
            print("💡 请在声音设置中启用'立体声混音'或'What U Hear'")
        
        return loopback_devices
    
    def enable_stereo_mix_guide(self):
        """启用立体声混音的指导"""
        guide = """
        🔧 启用立体声混音的步骤：
        
        1. 右键点击系统托盘的音量图标
        2. 选择"打开声音设置"
        3. 点击"声音控制面板"
        4. 切换到"录制"选项卡
        5. 右键空白区域，选择"显示禁用的设备"
        6. 找到"立体声混音"或"Stereo Mix"
        7. 右键点击，选择"启用"
        8. 设为默认录制设备（可选）
        
        🔄 或者尝试以下方法：
        - 使用虚拟音频线缆（VB-Cable）
        - 使用第三方音频捕获软件
        """
        print(guide)
    
    def capture_system_audio(self, device_index=None, callback=None):
        """捕获系统音频"""
        self.callback_func = callback
        
        if device_index is None:
            # 尝试找到默认的环回设备
            loopback_devices = self.list_loopback_devices()
            if not loopback_devices:
                self.enable_stereo_mix_guide()
                return False
            device_index = loopback_devices[0][0]
        
        try:
            # 创建音频流用于捕获系统播放的声音
            self.stream = self.audio.open(
                format=pyaudio.paFloat32,
                channels=self.channels,
                rate=self.sample_rate,
                input=True,
                input_device_index=device_index,
                frames_per_buffer=self.chunk_size,
                stream_callback=self._audio_callback
            )
            
            self.is_recording = True
            self.stream.start_stream()
            
            print(f"✅ 开始捕获系统音频，设备索引: {device_index}")
            return True
            
        except Exception as e:
            print(f"❌ 捕获系统音频失败: {e}")
            return False
    
    def _audio_callback(self, in_data, frame_count, time_info, status):
        """音频回调函数"""
        if self.is_recording:
            try:
                # 转换音频数据
                audio_data = np.frombuffer(in_data, dtype=np.float32)
                
                # 计算音量
                if self.channels == 2:
                    # 立体声转单声道
                    mono_data = audio_data.reshape(-1, 2).mean(axis=1)
                else:
                    mono_data = audio_data
                
                volume = np.sqrt(np.mean(mono_data**2))
                self.audio_buffer.append((time.time(), volume, audio_data.copy()))
                
                self.volume=volume
                # 检查音量阈值
                if volume >= self.volume_threshold and self.callback_func:
                    self.callback_func(volume, audio_data)
                
            except Exception as e:
                print(f"音频处理错误: {e}")
        
        return (in_data, pyaudio.paContinue)
    
    def stop_capture(self):
        """停止捕获"""
        self.is_recording = False
        if self.stream:
            self.stream.stop_stream()
            self.stream.close()
            self.stream = None
        print("✅ 音频捕获已停止")
    
    def save_captured_audio(self, filename, duration_seconds=10):
        """保存捕获的音频"""
        if not self.audio_buffer:
            print("❌ 没有音频数据可保存")
            return False
        
        try:
            # 提取最近的音频数据
            current_time = time.time()
            recent_audio = []
            
            for timestamp, volume, audio_data in self.audio_buffer:
                if current_time - timestamp <= duration_seconds:
                    recent_audio.append(audio_data)
            
            if not recent_audio:
                print("❌ 没有足够的音频数据")
                return False
            
            # 合并音频数据
            combined_audio = np.concatenate(recent_audio)
            
            # 转换为16位整数
            audio_int16 = (combined_audio * 32767).astype(np.int16)
            
            # 保存为WAV文件
            with wave.open(filename, 'wb') as wav_file:
                wav_file.setnchannels(self.channels)
                wav_file.setsampwidth(2)  # 16位 = 2字节
                wav_file.setframerate(self.sample_rate)
                wav_file.writeframes(audio_int16.tobytes())
            
            print(f"✅ 音频已保存到: {filename}")
            return True
            
        except Exception as e:
            print(f"❌ 保存音频失败: {e}")
            return False
    
    def get_current_volume(self):
        """获取当前音量"""
        # if self.audio_buffer:
        #     return self.audio_buffer[-1][1]
        return self.volume
    
    def __del__(self):
        """清理资源"""
        self.stop_capture()
        if self.audio:
            self.audio.terminate()

# 使用示例
if __name__ == "__main__":
    def audio_detected(volume, audio_data):
        print(f"🔊 检测到音频: 音量 {volume:.4f}")
    
    capture = ApplicationSoundCapture()
    capture.volume_threshold = 0.1  # 设置音量阈值
    
    # 列出可用设备
    devices = capture.list_loopback_devices()
    
    if devices:
        # 使用第一个找到的环回设备
        device_index = devices[0][0]
        
        if capture.capture_system_audio(device_index, audio_detected):
            try:
                print("正在捕获系统音频... 按 Ctrl+C 停止")
                while True:
                    volume = capture.get_current_volume()
                    print(f"当前音量: {volume:.4f}", end='\r')
                    time.sleep(0.1)
                    
            except KeyboardInterrupt:
                print("\n停止捕获...")
                capture.stop_capture()
                
                # 保存最后10秒的音频
                capture.save_captured_audio("captured_audio.wav", 10)
        else:
            print("❌ 音频捕获启动失败")
    else:
        print("❌ 没有找到可用的环回设备")